rm(list=ls())

setwd("/Users/chuangchen/Library/CloudStorage/OneDrive-UniversityofPittsburgh/Pela project")

library(haven)
library(dplyr)
library(ggplot2)
library(ggpubr)
options(scipen=999)

#######################################
# plots for model 1

#install.packages("ggpubr")
library(ggpubr)
library(haven)
library(dplyr)
library(ggplot2)


######################################################################
# test contextual effects
# wide figure
left_right <- c("left", "right")
left_right_C <- c("Left", "Right")
context_pela <- c("vola", "polar_pela", "ENP")

context <- c("client", "exclu", "psi",
             "pid_h", "polar", "poor", 
             "reli", "stable","cohesion", 
             "high_crime", "high_equal",
             "high_edu_gini")
context_list_h_l <- c("More Clientelist", "Less Clientelist",
                      "High Poli Exclu", "Low Poli Exclu",
                      "Low PSI","High PSI", 
                      "High PID", "Low PID",
                      "High Polarization", "Low Polarization",
                      "Poor Economy", "Good Economy",
                      "High P-System Reli", "Low P-System Reli",
                      "High Poli Stability", "Low Poli Stability",
                      "High Legis P Cohe", "Low Legis P Cohe",
                      "High Crime", "Low Crime",
                      "High Equality", "Low Equality",
                      "High Education Gini", "Low Education Gini")
context_list_h_l_pela <- c("High Electoral Vola", "Low Electoral Vola",
                           "High Polarization", "Low Polarization",
                           "High ENP", "Low ENP")
# merge context and context_list_h_l to a data frame
context_df <- data.frame(context=rep(context, each=2), context_list_h_l=context_list_h_l)
context_df_pela <- data.frame(context=rep(context_pela, each=2), context_list_h_l=context_list_h_l_pela)
                      
label <- c("Clientelism", "Political Exclusion", "Low PSI", 
           "High PID", "Polarization", "Poor Economy", 
           "Party-System Religion", "Political Stability",
           "Legislative Party Cohesion", "High Crime", "High Equality",
           "High Education Gini")
label_pela <- c("Electoral Volatility", "Polarization", "ENP")

df <- data.frame(context=rep(context, each=2), label=rep(label, each=2))
# create a new column in df, for each element in left_right, fill it with the corresponding value in context
df$left_right <- rep(left_right, times=12)
df$left_right_C <- rep(left_right_C, times=12)

# a new column with data name
df$data <- paste0("model1_pred_prob_", df$left_right, "_", df$context, "_inter_effect", ".dta")
df$plot_title <- paste0("model1_pred_prob_", df$context, "_inter_effect", ".pdf")

# a new column with y axis title
df$y_title <- paste0("Effect of ", df$label, " on Being ", df$left_right_C)


df_pela <- data.frame(context=rep(context_pela, each=2), label=rep(label_pela, each=2))
df_pela$left_right <- rep(left_right, times=3)
df_pela$left_right_C <- rep(left_right_C, times=3)
df_pela$data <- paste0("model1_pred_prob_", df_pela$left_right, "_", df_pela$context, "_inter_effect", ".dta")
df_pela$plot_title <- paste0("model1_pred_prob_", df_pela$context, "_inter_effect", ".pdf")


df_append <- data.frame()
for(i in 1:nrow(df)){
  data1 <- read_dta(df[i, "data"]) %>% select('_at1', "_at2", '_margin')
  
  data1$eco <- rep(c("conservative(Ec)","middle(Em)","progressive(Ep)",
                     "conservative(Ec)","middle(Em)","progressive(Ep)",
                     "conservative(Ec)","middle(Em)","progressive(Ep)"), 
                   each=2)
  data1$eco_brief <- rep(c("Ec","Em","Ep",
                     "Ec","Em","Ep",
                     "Ec","Em","Ep"), each=2)
  data1$`Moral Position` <- rep(c("Mc","Mc","Mc",
                                  "Mm","Mm","Mm",
                                  "Mp","Mp","Mp"), each=2)
  colnames(data1) <- c("eco_moral", "context", "pred", "eco", "eco_brief", "Moral Position")
  # group by eco_moral, create client1 is the value of pred when client = 1, and client0 is the value of pred when client = 0
  data1 <- data1 %>% group_by(eco_moral) %>% 
    mutate(context1 = pred[context == 1], 
           context0 = pred[context == 0]) %>% ungroup()
  # group by eco_mora, only keep one row in each group
  data1 <- data1 %>% group_by(eco_moral) %>% slice(1) %>% ungroup()
  # drop the client column
  data1 <- data1 %>% select(-context)
  
  data1$effect <- data1$context1 - data1$context0
  data1$context <- df[i, "label"]
  data1$left_right <- df[i, "left_right_C"]
  # append to df_append
  df_append <- rbind(df_append, data1)
}


df_append_pela <- data.frame()
for(i in 1:nrow(df_pela)){
  data1 <- read_dta(df_pela[i, "data"]) %>% select('_at1', "_at2", '_margin')
  
  data1$eco <- rep(c("conservative(Ec)","middle(Em)","progressive(Ep)",
                     "conservative(Ec)","middle(Em)","progressive(Ep)",
                     "conservative(Ec)","middle(Em)","progressive(Ep)"), 
                   each=2)
  data1$eco_brief <- rep(c("Ec","Em","Ep",
                           "Ec","Em","Ep",
                           "Ec","Em","Ep"), each=2)
  data1$`Moral Position` <- rep(c("Mc","Mc","Mc",
                                  "Mm","Mm","Mm",
                                  "Mp","Mp","Mp"), each=2)
  colnames(data1) <- c("eco_moral", "context", "pred", "eco", "eco_brief", "Moral Position")
  # group by eco_moral, create client1 is the value of pred when client = 1, and client0 is the value of pred when client = 0
  data1 <- data1 %>% group_by(eco_moral) %>% mutate(context1 = pred[context == 1], 
                                                    context0 = pred[context == 0]) %>% ungroup()
  # group by eco_mora, only keep one row in each group
  data1 <- data1 %>% group_by(eco_moral) %>% slice(1) %>% ungroup()
  # drop the client column
  data1 <- data1 %>% select(-context)
  
  data1$effect <- data1$context1 - data1$context0
  data1$context <- df_pela[i, "label"]
  data1$left_right <- df_pela[i, "left_right_C"]
  # append to df_append
  df_append_pela <- rbind(df_append_pela, data1)
}




plot1 <- ggplot(df_append, aes(x=eco_brief)) + 
  geom_vline(aes(xintercept="Ec"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Em"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Ep"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_point(aes(y=effect, shape=`Moral Position`, color=`Moral Position`), 
             position=position_dodge(0.5), size=2) +
  ggtitle(df[i,"Effect of Context"]) +
  # add a red dotted line at y = 0
  geom_hline(yintercept=0, linetype="dashed", color = "red") +
  # put the tile at the center
  theme(plot.title = element_text(hjust = 0.5)) +
  theme(panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        panel.background = element_blank(), 
        axis.line = element_line(colour = "black"),
        axis.text = element_text(size = 12),
        axis.title.y = element_text(size = 14),
        axis.title.x = element_text(size = 14),
        axis.text.x = element_text(color = "black"),
        axis.text.y = element_text(color = "black")) +
  ylab("Contextual Effect") +
  xlab("Economic Position") +
  scale_x_discrete(breaks=c("Ec", "Em", "Ep")) +
  scale_y_continuous(limits = c(-0.3, 0.2), 
                     labels = scales::number_format(scale = 1)) +
  scale_color_manual(values = c("Mc"="#FF99FF", 
                                "Mm"="chartreuse3", 
                                "Mp"="#0099FF"))+
  guides(color=guide_legend(title="Moral Position")) + 
  theme(legend.position="bottom",
        legend.text = element_text(size=10)) +
  # divide by left_right and context
  facet_grid(left_right ~ context, scales="free_y") 

ggsave(filename = "all_context.png", plot1, width=15, height=9)





plot1 <- ggplot(df_append_pela, aes(x=eco_brief)) + 
  geom_vline(aes(xintercept="Ec"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Em"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Ep"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_point(aes(y=effect, shape=`Moral Position`, color=`Moral Position`), 
             position=position_dodge(0.5), size=2) +
  ggtitle(df_pela[i,"Effect of Context"]) +
  # add a red dotted line at y = 0
  geom_hline(yintercept=0, linetype="dashed", color = "red") +
  # put the tile at the center
  theme(plot.title = element_text(hjust = 0.5)) +
  theme(panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        panel.background = element_blank(), 
        axis.line = element_line(colour = "black"),
        axis.text = element_text(size = 12),
        axis.title.y = element_text(size = 14),
        axis.title.x = element_text(size = 14),
        axis.text.x = element_text(color = "black"),
        axis.text.y = element_text(color = "black")) +
  ylab("Contextual Effect") +
  xlab("Economic Position") +
  scale_x_discrete(breaks=c("Ec", "Em", "Ep")) +
  scale_y_continuous(limits = c(-0.3, 0.2), 
                     labels = scales::number_format(scale = 1)) +
  scale_color_manual(values = c("Mc"="#FF99FF", 
                                "Mm"="chartreuse3", 
                                "Mp"="#0099FF"))+
  guides(color=guide_legend(title="Moral Position")) + 
  theme(legend.position="bottom",
        legend.text = element_text(size=10)) +
  # divide by left_right and context
  facet_grid(left_right ~ context, scales="free_y") 

ggsave(filename = "all_context_pela.png", plot1, width=6, height=8)













plot_context <- function(df, i, context_df){
  data1 <- read_dta(df[i, "data"]) %>% select('_at1', "_at2", '_margin', 
                                              '_ci_lb', '_ci_ub')
  data1$eco_brief <- rep(c("Ec","Em","Ep",
                           "Ec","Em","Ep",
                           "Ec","Em","Ep"), each=2)
  data1$`moral` <- rep(c("Mc","Mc","Mc",
                         "Mm","Mm","Mm",
                         "Mp","Mp","Mp"), each=2)
  colnames(data1) <- c("eco_moral", "context", "pred", "ll", "ul", "eco_brief", "moral")
  # extract rows with context=0 to data2, only keep context=1 in data1
  data2 <- data1 %>% filter(context == 0)
  data1 <- data1 %>% filter(context == 1)
  data1$context <- data1$context %>% as.factor()
  data2$context <- data2$context %>% as.factor()
  
  # extract the value of context_df$context_list_h_l when context_df$context = df[i,"context"]
  context_short <- df[i, "context"]
  context_long <- context_df$context_list_h_l[context_df$context == context_short]
  
  plot1 <- ggplot() + 
    geom_vline(data=data1, aes(xintercept="Ec"), linewidth=25, color = alpha("gray", 0.2)) +
    geom_vline(data=data1, aes(xintercept="Em"), linewidth=25, color = alpha("gray", 0.2)) +
    geom_vline(data=data1, aes(xintercept="Ep"), linewidth=25, color = alpha("gray", 0.2)) +
    geom_pointrange(data=data1, aes(x=as.numeric(factor(eco_brief))-0.08, 
                                    color = moral,
                                    y=pred, ymin=ll, ymax=ul,shape=context),
                    position = position_dodge(0.5), 
                    alpha=1) +
    geom_pointrange(data=data2, aes(x=eco_brief, 
                                    color = moral,
                                    y=pred, ymin=ll, ymax=ul,shape=context),
                    position = position_dodge(0.5), 
                    alpha=0.4) +
    ggtitle(df[i, "label"]) +
    ylab(paste0("Predicted Probability of Being ", df[i, "left_right_C"])) +
    xlab("Economic Position") +
    scale_x_discrete(breaks=c("Ec", "Em", "Ep"),
                     labels=c("Ec", "Em", "Ep")) +
    scale_y_continuous(limits = c(-0.1,1), labels = scales::percent) +
    scale_color_manual(values = c("Mc"="#FF99FF", "Mm"="chartreuse3", "Mp"="#0099FF"),
                       labels = c("Mc", "Mm", "Mp")) +
    scale_shape_manual(values = c("1"=16,"0"=1),
                       labels = c("1"=context_long[1], "0"=context_long[2])) +
    guides(color=guide_legend(title="Moral Position"),
           shape=guide_legend(title="Context", reverse = FALSE, order=1)) +
    theme(legend.position="bottom",
          legend.text = element_text(size=10),
          legend.box = "vertical",
          legend.margin=margin(t=0, r=0, b=0, l=0),
          plot.title = element_text(hjust = 0.5),
          panel.grid.major = element_blank(), 
          panel.grid.minor = element_blank(),
          panel.background = element_blank(), 
          axis.line = element_line(colour = "black"),
          axis.text = element_text(size = 12),
          axis.title.y = element_text(size = 14),
          axis.title.x = element_text(size = 14),
          axis.text.x = element_text(color = "black"),
          axis.text.y = element_text(color = "black"))
  
  
}



for(i in 1:nrow(df)){
 plot1 <- plot_context(df, i, context_df)
 assign(paste0("plot_", df[i,"context"], "_", df[i, "left_right"]), plot1)
}

for(i in 1:nrow(df_pela)){
  plot1 <- plot_context(df_pela, i, context_df_pela)
  assign(paste0("plot_", df_pela[i,"context"], "_", df[i, "left_right"]), plot1)
}


# first use ggarrange to put all the plots of left in one page
# do not share legend of context
plot_left <- ggarrange(plot_client_left, plot_exclu_left, plot_psi_left, 
                       plot_pid_h_left, plot_polar_left, plot_poor_left,
                       plot_reli_left, plot_stable_left, plot_cohesion_left, 
                       plot_high_crime_left, plot_high_equal_left,plot_high_edu_gini_left,
                       ncol=3,nrow=4,common.legend = FALSE)
ggsave("plot_left.png", plot_left, width=15, height=20, units="in", dpi=300)

# then use ggarrange to put all the plots of right in one page
# do not share legend of context
plot_right <- ggarrange(plot_client_right, plot_exclu_right, plot_psi_right, 
                        plot_pid_h_right, plot_polar_right, plot_poor_right,
                        plot_reli_right, plot_stable_right, plot_cohesion_right, 
                        plot_high_crime_right, plot_high_equal_right,plot_high_edu_gini_right,
                        ncol=3, nrow=4, common.legend = FALSE)
ggsave("plot_right.png", plot_right, width=15, height=20, units="in", dpi=300)


# plot pela
plot_pela <- ggarrange(plot_vola_left, plot_ENP_left, plot_polar_pela_left,
                       plot_vola_right, plot_ENP_right, plot_polar_pela_right,
                       ncol=3, nrow=2, common.legend = FALSE)
ggsave("plot_pela.png", plot_pela, width=13, height=9, units="in", dpi=300)

# legend at the bottom
client_left_right <- ggarrange(plot_client_left, plot_client_right, ncol=2, nrow=1, common.legend = TRUE,
                               legend = "bottom")
ggsave("client_left_right.png", client_left_right, width=10, height=5, units="in", dpi=300)

#################################################################################
library(tidyr)
#test
# generate a new variable MmEm, the value of MmEm is the value of pred 
df_append_pela$MmEm_C1 <- df_append_pela$context1
df_append_pela$MmEm_C0 <- df_append_pela$context0
# if eco_moral is not 5, then MmEm is NA
df_append_pela$MmEm_C1[df_append_pela$eco_moral != 5] <- NA
df_append_pela$MmEm_C0[df_append_pela$eco_moral != 5] <- NA
# within each context, within each left_right, if MmEm is NA, replace it with the value that is not NA
df_append_pela <- df_append_pela %>% group_by(context, left_right) %>% fill(MmEm_C1, .direction = "updown")
df_append_pela <- df_append_pela %>% group_by(context, left_right) %>% fill(MmEm_C0, .direction = "updown")

df_append_pela$diff_to_MmEm_C1 <- df_append_pela$context1 - df_append_pela$MmEm_C1
df_append_pela$diff_to_MmEm_C0 <- df_append_pela$context0 - df_append_pela$MmEm_C0

df_append_pela$diff_context <- df_append_pela$diff_to_MmEm_C1 - df_append_pela$diff_to_MmEm_C0




plot1 <- ggplot(df_append_pela, aes(x=eco_brief)) + 
  geom_vline(aes(xintercept="Ec"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Em"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Ep"), linewidth=13, color = alpha("gray", 0.2)) +
  geom_point(aes(y=diff_context, shape=`Moral Position`, color=`Moral Position`), 
             position=position_dodge(0.5), size=2) +
  ggtitle(df[i,"Effect of Context"]) +
  # add a red dotted line at y = 0
  geom_hline(yintercept=0, linetype="dashed", color = "red") +
  # put the tile at the center
  theme(plot.title = element_text(hjust = 0.5)) +
  theme(panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        panel.background = element_blank(), 
        axis.line = element_line(colour = "black"),
        axis.text = element_text(size = 12),
        axis.title.y = element_text(size = 14),
        axis.title.x = element_text(size = 14),
        axis.text.x = element_text(color = "black"),
        axis.text.y = element_text(color = "black")) +
  ylab("Contextual Effect") +
  xlab("Economic Position") +
  scale_x_discrete(breaks=c("Ec", "Em", "Ep")) +
  scale_y_continuous(limits = c(-0.3, 0.2), 
                     labels = scales::number_format(scale = 1)) +
  scale_color_manual(values = c("Mc"="#FF99FF", 
                                "Mm"="chartreuse3", 
                                "Mp"="#0099FF"))+
  guides(color=guide_legend(title="Moral Position")) + 
  theme(legend.position="bottom",
        legend.text = element_text(size=10)) +
  # divide by left_right and context
  facet_grid(left_right ~ context, scales="free_y") 




df_append_pela$C1 <- "1"
df_append_pela$C0 <- "0"

plot <- ggplot(df_append_pela) + 
  geom_vline(aes(xintercept="Ec"), linewidth=25, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Em"), linewidth=25, color = alpha("gray", 0.2)) +
  geom_vline(aes(xintercept="Ep"), linewidth=25, color = alpha("gray", 0.2)) +
  geom_hline(aes(yintercept=0), linetype="dashed", color = "red") +
  geom_point(data=df_append_pela, aes(x=as.numeric(factor(eco_brief))-0.08, 
                                  color = `Moral Position`,
                                  y=diff_to_MmEm_C1),
                  position = position_dodge(0.5), 
                  alpha=1) +
  geom_point(data=df_append_pela, aes(x=eco_brief, 
                                  color = `Moral Position`,
                                  y=diff_to_MmEm_C0),
                  position = position_dodge(0.5), 
                  alpha=0.3) +
  ggtitle(df[i, "label"]) +
  xlab("Economic Position") +
  scale_x_discrete(breaks=c("Ec", "Em", "Ep"),
                   labels=c("Ec", "Em", "Ep")) +
  scale_y_continuous(limits = c(-0.5,0.5), labels = scales::percent) +
  scale_color_manual(values = c("Mc"="#FF99FF", "Mm"="chartreuse3", "Mp"="#0099FF"),
                     labels = c("Mc", "Mm", "Mp")) +
  guides(color=guide_legend(title="Moral Position"),
         shape=guide_legend(title="Context", reverse = FALSE, order=1)) +
  theme(legend.position="bottom",
        legend.text = element_text(size=10),
        legend.box = "vertical",
        legend.margin=margin(t=0, r=0, b=0, l=0),
        plot.title = element_text(hjust = 0.5),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        panel.background = element_blank(), 
        axis.line = element_line(colour = "black"),
        axis.text = element_text(size = 12),
        axis.title.y = element_text(size = 14),
        axis.title.x = element_text(size = 14),
        axis.text.x = element_text(color = "black"),
        axis.text.y = element_text(color = "black")) +
  facet_grid(left_right ~ context, scales="free_y")


ggsave(filename = "", plot1, width=15, height=9)
